package framework.storage;

import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.model.*;
import lombok.Getter;
import lombok.Setter;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

public class FileStorageS3OutputStream extends OutputStream {

    // 最小分块大小
    private final long PART_SIZE;

    private final AmazonS3 s3Client;
    private final String bucketName;
    private final String key;
    private final ObjectMetadata metadata;

    private String uploadId;         // 分块上传ID
    private List<PartETag> partETags;// 分块上传标签列表
    @Getter
    @Setter
    private boolean writeCompleted;

    public FileStorageS3OutputStream(AmazonS3 s3Client, String bucketName, String key, ObjectMetadata metadata, long partSize) {
        this.PART_SIZE = partSize;
        this.s3Client = s3Client;
        this.bucketName = bucketName;
        this.key = key;
        this.metadata = metadata;

        this.partETags = new ArrayList<>();
        this.writeCompleted = false;
    }

    /**
     * 写入单个字节
     */
    @Override
    public void write(int b) throws IOException {
        throw new UnsupportedOperationException("Not support method");
    }

    /**
     * 写入字节数组
     */
    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        if (len == 0) {
            return;
        }

        // 首次调用
        if (partETags.size() == 0) {
            if (len < PART_SIZE) {
                putObject(b, off, len);
                return;
            }
            InitiateMultipartUploadRequest initRequest = new InitiateMultipartUploadRequest(bucketName, key, metadata);
            InitiateMultipartUploadResult initResponse = s3Client.initiateMultipartUpload(initRequest);
            uploadId = initResponse.getUploadId();
        }

        // 分块上传
        try (ByteArrayInputStream inputStream = new ByteArrayInputStream(b, off, len)) {
            UploadPartRequest uploadRequest = new UploadPartRequest()
                    .withBucketName(bucketName)
                    .withKey(key)
                    .withUploadId(uploadId)
                    .withPartNumber(partETags.size() + 1)
                    .withInputStream(inputStream)
                    .withPartSize(len);
            UploadPartResult result = s3Client.uploadPart(uploadRequest);
            partETags.add(result.getPartETag());
        }
    }

    @Override
    public void flush() throws IOException {
        // 底层处理逻辑，无需flush
    }

    private void putObject(byte[] b, int off, int len) throws IOException {
        // 普通上传模式：一次性上传所有数据
        try (ByteArrayInputStream inputStream = new ByteArrayInputStream(b, off, len)) {
            metadata.setContentLength(len);
            s3Client.putObject(new PutObjectRequest(bucketName, key, inputStream, metadata));
        }
    }

    /**
     * 关闭流并完成上传
     */
    @Override
    public void close() throws IOException {
        if (uploadId != null) {
            if (partETags.size() > 0 && isWriteCompleted()) {
                // 完成分块上传
                completeMultipartUpload();
                uploadId = null;
            } else {
                // 中止未写完的分块上传
                abortMultipartUpload();
                uploadId = null;
                throw new IOException("Upload aborted: Incomplete data write (not all content was processed) or setWriteCompleted for copy completion was not invoked");
            }
        }
    }

    /**
     * 完成分块上传
     */
    private void completeMultipartUpload() {
        CompleteMultipartUploadRequest uploadRequest = new CompleteMultipartUploadRequest(bucketName, key, uploadId, partETags);
        s3Client.completeMultipartUpload(uploadRequest);
    }

    /**
     * 中止分块上传
     */
    private void abortMultipartUpload() {
        AbortMultipartUploadRequest uploadRequest = new AbortMultipartUploadRequest(bucketName, key, uploadId);
        s3Client.abortMultipartUpload(uploadRequest);
    }
}
